import pandas as pd
from collections import defaultdict

def get_core_task(question):
    # 编程类
    if any(word in question.lower() for word in ['编程', 'java', 'python', 'debug', '代码', '程序', 'd3']):
        return '编程'
    # 报告/总结类
    elif any(word in question for word in ['报告', '总结']):
        return '报告'
    # 论文相关
    elif any(word in question for word in ['论文', '文献']):
        return '论文'
    # 作业类
    elif '作业' in question:
        return '作业'
    # 翻译类
    elif '翻译' in question or '英文' in question:
        return '翻译'
    # PPT
    elif 'ppt' in question.lower():
        return 'PPT'
    # 写作类
    elif any(word in question for word in ['作文', '文章', '文稿', '感悟']):
        return '写作'
    # 查找/查询
    elif any(word in question for word in ['查找', '查题', '找一个']):
        return '查找'
    # 提炼/总结
    elif any(word in question for word in ['提炼', '概括']):
        return '提炼'
    # 问答类
    elif any(word in question for word in ['是什么', '什么意思', '这道题', '问题', '想法']):
        return '问答'
    # 思路/框架/规划
    elif any(word in question for word in ['思路', '框架', '大纲', '规划', '计划']):
        return '思路'
    # 改写/优化
    elif any(word in question for word in ['改写', '优化', '降重']):
        return '改写'
    # 分析
    elif any(word in question for word in ['分析', '研究', '物性']):
        return '分析'
    # 数学
    elif any(word in question for word in ['高数', '数学']):
        return '数学'
    # 图像处理
    elif any(word in question for word in ['图片', '降噪']):
        return '图像'
    # 生活类
    elif '生活' in question:
        return '生活'
    
    return '待分类'

# 读取CSV文件
df = pd.read_csv('src/assets/processed.csv')

# 提取非空的specific_question和对应的major
data = df[['specific_question', 'major']].dropna(subset=['specific_question'])

# 创建结果DataFrame
results = []

# 处理数据
for _, row in data.iterrows():
    task = get_core_task(row['specific_question'])
    results.append({
        'major': row['major'],
        'specific_question': row['specific_question'],
        'category': task
    })

# 转换为DataFrame并保存
results_df = pd.DataFrame(results)

# 保存为CSV文件
results_df.to_csv('src/assets/categorized_questions.csv', index=False)

# 打印统计结果
print("\n=== 分类统计 ===")
category_counts = results_df['category'].value_counts()
for category, count in category_counts.items():
    print(f"{category}: {count}次")

# 打印未分类的问题
unclassified = results_df[results_df['category'] == '待分类']
if not unclassified.empty:
    print("\n=== 待分类问题 ===")
    for _, row in unclassified.iterrows():
        print(f"{row['major']}: {row['specific_question']}")